#include <stdio.h>
#include <pthread.h>
#include <stdlib.h>
#include <unistd.h>
#include <queue>
#include <unordered_map>
#include <cmath>
#include <dirent.h>
#include <string.h>
#include "opencv2/opencv.hpp"
#include "opencv2/imgproc/types_c.h"
#include "opencv2/highgui/highgui.hpp"
#include "acllite_dvpp_lite/ImageProc.h"
#include "acllite_om_execute/ModelProc.h"
#include "acllite_dvpp_lite/VideoRead.h"
#include "acllite_common/Queue.h"
#include <sys/time.h>

using namespace std;
using namespace acllite;
using namespace cv;
namespace {
    aclrtContext context = nullptr;
    // Set the number of channels to open the corresponding number of video streams
    const int channelNum = 16;
    // Set showNum for cv::imshow windows arrangement rules
    const int showNum = 4;
    const uint32_t modelWidth = 640;
    const uint32_t modelHeight = 640;

    const uint32_t fullScreenWidth = 1920;
    const uint32_t fullScreenHeight = 1050;
    const uint32_t titleHeight = 30;
    const int channel[36] = {0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,
        17,18,19,20,21,22,23,24,25,26,27,28,29,30,31,32,33,34,35};
    bool exitFlag = false;

    struct Config {
        int32_t deviceId;
        int channel;
    };

    struct MsgData {
        std::shared_ptr<uint8_t> data = nullptr;
        uint32_t size = 0;
        uint32_t width = 0;
        uint32_t height = 0;
        int ChannelId = 0;
        bool videoEnd = false;
        cv::Mat srcImg;
    };

    struct MsgOut {
        cv::Mat srcImg;
        int ChannelId = 0;
        bool videoEnd = false;
        vector<InferenceOutput> inferOutputs;
    };

    Queue<MsgData> msgDataQueue(32);
    Queue<MsgOut> msgOutQueue(32);
    size_t imageInfoSize;
    void* imageInfoBuf;
    std::vector<string> winNameList;
    std::string displayMode = "";
    const static std::vector<std::string> yolov3Label = { "person", "bicycle", "car", "motorbike",
        "aeroplane", "bus", "train", "truck", "boat",
        "traffic light", "fire hydrant", "stop sign", "parking meter",
        "bench", "bird", "cat", "dog", "horse",
        "sheep", "cow", "elephant", "bear", "zebra",
        "giraffe", "backpack", "umbrella", "handbag", "tie",
        "suitcase", "frisbee", "skis", "snowboard", "sports ball",
        "kite", "baseball bat", "baseball glove", "skateboard", "surfboard",
        "tennis racket", "bottle", "wine glass", "cup",
        "fork", "knife", "spoon", "bowl", "banana",
        "apple", "sandwich", "orange", "broccoli", "carrot",
        "hot dog", "pizza", "donut", "cake", "chair",
        "sofa", "potted plant", "bed", "dining table", "toilet",
        "TV monitor", "laptop", "mouse", "remote", "keyboard",
        "cell phone", "microwave", "oven", "toaster", "sink",
        "refrigerator", "book", "clock", "vase", "scissors",
        "teddy bear", "hair drier", "toothbrush" };

    enum BBoxIndex { TOPLEFTX = 0, TOPLEFTY, BOTTOMRIGHTX, BOTTOMRIGHTY, SCORE, LABEL };
    // bounding box line solid
    const uint32_t lineSolid = 2;
    // opencv draw label params.
    const double fountScale = 0.5;
    const cv::Scalar fontColor(0, 0, 255);
    const uint32_t labelOffset = 11;
    // opencv color list for boundingbox
    const vector<cv::Scalar> colors {
        cv::Scalar(237, 149, 100), cv::Scalar(0, 215, 255), cv::Scalar(50, 205, 50),
        cv::Scalar(139, 85, 26) };
    struct Rect {
        uint32_t ltX = 0;
        uint32_t ltY = 0;
        uint32_t rbX = 0;
        uint32_t rbY = 0;
    };
    struct BBox {
        Rect rect;
        uint32_t score;
        string text;
    };

}

bool Initparam(int argc, char *argv[])
{
    DIR *dir;
    int outputDisplay = 1;
    if ((dir = opendir("./out")) == NULL)
        system("mkdir ./out");
    int paramNum = 2;
    if(argc != paramNum) {
        LOG_PRINT("[ERROR] please choose output display mode: [./main imshow] [./main stdout]");
        return false;
    }
    displayMode = argv[outputDisplay];
    return true;
}

void Draw(cv::Mat &srcImg, int channel) {
    cv::Mat resizedImg;
    cv::resize(srcImg, resizedImg, cv::Size(fullScreenWidth/showNum, (fullScreenHeight-titleHeight)/showNum));
    cv::imshow(winNameList[channel], resizedImg);
    cv::waitKey(1);
}

void GetResult(std::vector<InferenceOutput>& inferOutputs,
    cv::Mat& srcImage, uint32_t modelWidth, uint32_t modelHeight,
    int channelId)
{
    float *detectData = static_cast<float *>(inferOutputs[0].data.get());
    uint32_t* boxNum = static_cast<uint32_t *>(inferOutputs[1].data.get());
    if (boxNum == nullptr) return;

    uint32_t totalBox = boxNum[0];
    int srcWidth = srcImage.cols;
    int srcHeight = srcImage.rows;
    float widthScale = (float)(srcWidth) / modelWidth;
    float heightScale = (float)(srcHeight) / modelHeight;
    float finalScale = (widthScale > heightScale) ? widthScale : heightScale;
    string textPrint = "[";
    for (uint32_t i = 0; i < totalBox; i++) {
        cv::Point p1, p2;
        uint32_t score = uint32_t(detectData[totalBox * SCORE + i] * 100);
        if (score < 70) {
            continue;
        }
        p1.x = detectData[totalBox * TOPLEFTX + i] * finalScale;
        p1.y = detectData[totalBox * TOPLEFTY + i] * finalScale;
        p2.x = detectData[totalBox * BOTTOMRIGHTX + i] * finalScale;
        p2.y = detectData[totalBox * BOTTOMRIGHTY + i] * finalScale;
        uint32_t objIndex = (uint32_t)detectData[totalBox * LABEL + i];
        string text = yolov3Label[objIndex] + std::to_string(score) + "\%";
        textPrint = textPrint + text + " ";
        cv::rectangle(srcImage, p1, p2, colors[i % colors.size()], lineSolid);
        cv::putText(srcImage, text, cv::Point(p1.x, p1.y + labelOffset),
                    cv::FONT_HERSHEY_COMPLEX, fountScale, fontColor);
    }
    textPrint = textPrint + "]";
    if (displayMode == "imshow") {
        Draw(srcImage, channelId);
    } else if (displayMode == "stdout") {
        if (textPrint == "[]") {
            cout<<"Channel-" << channelId << "-result: "<<"none object detected."<<endl;
        } else {
            cout<<"Channel-" << channelId << "-result: "<< textPrint <<endl;
        }
    } else {
        LOG_PRINT("[ERROR] output display mode not supported.");
    }
    return;
}

void* GetInput(void* arg) {
    bool ret = SetCurContext(context);
    CHECK_RET(ret, LOG_PRINT("[ERROR] set cur context for pthread  %ld failed.", pthread_self()); return NULL);
    Config cfg = *(Config *)arg;
    string videoPath = "../data/test.mp4";
    VideoRead cap(videoPath, cfg.deviceId);
    CHECK_RET(cap.IsOpened(), LOG_PRINT("[ERROR] open %s failed.", videoPath.c_str()); return NULL);
    ImageProc imageProcess;
    ImageData frame;
    ImageSize modelSize(modelWidth, modelHeight);
    LOG_PRINT("[INFO] start to decode...");
    uint32_t frameRate =30;
    uint32_t cnt = 1;
    int isHost = GetRunMode();
    while(1) {
        ret = cap.Read(frame);
        if (ret) {
            if (cnt % frameRate == 0) {
                ImageData dst;
                imageProcess.Resize(frame, dst, modelSize, RESIZE_PROPORTIONAL_UPPER_LEFT);
                MsgData msgData;
                msgData.data = dst.data;
                msgData.size = dst.size;
                msgData.width = dst.width;
                msgData.height = dst.height;
                msgData.ChannelId = cfg.channel;
                msgData.videoEnd = false;
                cv::Mat yuv_img(frame.height*1.5, frame.width, CV_8UC1);
                if (isHost) {
                    void* hostDataBuffer = CopyDataToHost(frame.data.get(), frame.size);
                    memcpy(yuv_img.data, (unsigned char*)hostDataBuffer, frame.size);
                    FreeHostMem(hostDataBuffer);
                    hostDataBuffer = nullptr;
                } else {
                    memcpy(yuv_img.data, (unsigned char*)frame.data.get(), frame.size);
                }
                cv::cvtColor(yuv_img, msgData.srcImg, cv::COLOR_YUV2RGB_NV21);
                while (1) {
                    if (msgDataQueue.Push(msgData)) {
                        break;
                    }
                    usleep(100);
                }
            }
            cnt++;
        } else {
            LOG_PRINT("[INFO] frame read end.");
            break;
        }
    }
    cap.Release();
    MsgData msgData;
    msgData.videoEnd = true;
    msgData.ChannelId = cfg.channel;
    while (1) {
        if (msgDataQueue.Push(msgData)) {
            break;
        }
        usleep(100);
    }
    return NULL;
}
 
void* ModelExecute(void* arg) {
    bool ret = SetCurContext(context);
    CHECK_RET(ret, LOG_PRINT("[ERROR] set cur context for pthread  %ld failed.", pthread_self()); return NULL);
    ModelProc modelProcess;
    string modelPath = "../model/yolov5s_nms.om";
    ret = modelProcess.Load(modelPath);
    CHECK_RET(ret, LOG_PRINT("[ERROR] load model %s failed.", modelPath.c_str()); return NULL);
    int num = *(int *)arg;
    while(num) {
        if(!msgDataQueue.Empty()) {
            MsgData msgData = msgDataQueue.Pop();
            if (msgData.videoEnd) {
                num--;
            }
            else {    
                ret = modelProcess.CreateInput(static_cast<void*>(msgData.data.get()), msgData.size, imageInfoBuf, imageInfoSize);
                CHECK_RET(ret, LOG_PRINT("[ERROR] Create model input failed."); break);
                MsgOut msgOut;
                msgOut.srcImg = msgData.srcImg;
                msgOut.ChannelId = msgData.ChannelId;
                msgOut.videoEnd = msgData.videoEnd;
                modelProcess.Execute(msgOut.inferOutputs);
                CHECK_RET(ret, LOG_PRINT("[ERROR] model execute failed."); break);
                while (1) {
                    if (msgOutQueue.Push(msgOut)) {
                        break;
                    }
                    usleep(100);
                }
            }
        }
    }
    modelProcess.DestroyResource();
    MsgOut msgOut;
    msgOut.videoEnd = true;
    while (1) {
        if (msgOutQueue.Push(msgOut)) {
            break;
        }
        usleep(100);
    }
    return  NULL;
}

void* PostProcess(void* arg) {
    while(1) {
        if(!msgOutQueue.Empty()) {
            MsgOut msgOut = msgOutQueue.Pop();
            usleep(100);
            if (msgOut.videoEnd) {
                break;
            }
            GetResult(msgOut.inferOutputs, msgOut.srcImg, modelWidth, modelHeight, msgOut.ChannelId);
        }
    }
    exitFlag = true;
    LOG_PRINT("[INFO] *************** all video get done ***************");
    return  NULL;
}

bool CreateBuffer(uint32_t modelWidth, uint32_t modelHeight)
{
    const float imageInfo[4] = {(float)modelWidth, (float)modelHeight,
                                (float)modelWidth, (float)modelHeight};
    imageInfoSize = sizeof(imageInfo);
    imageInfoBuf = CopyDataToDevice((void*)imageInfo, imageInfoSize);
    if (imageInfoBuf == nullptr) {
        LOG_PRINT("[ERROR] Copy image info to device failed");
        return false;
    }
    return true;
}

int main(int argc, char *argv[]) {
    int32_t deviceId = 0;
    AclLiteResource aclResource(deviceId);
    bool ret = aclResource.Init();
    CHECK_RET(ret, LOG_PRINT("[ERROR] InitACLResource failed."); return 1);
    context = aclResource.GetContext();    
    ret = Initparam(argc, argv);
    CHECK_RET(ret, LOG_PRINT("[ERROR] Initparam failed."); return 1);
    ret = CreateBuffer(modelWidth, modelHeight);
    CHECK_RET(ret, LOG_PRINT("[ERROR] CreateBuffer failed."); return 1);

    std::vector<Config> cfg;
    cfg.resize(channelNum);
    pthread_t preTids[channelNum], exeTids, posTids;
    for(int i = 0; i < channelNum; i++) {
        cfg[i].deviceId = deviceId;
        cfg[i].channel = i;
        pthread_create(&preTids[i], NULL, GetInput, (void*)&cfg[i]);
    }

    if (displayMode == "imshow") {
        for (int i = channelNum - 1; i >= 0; i--) {
            std::string head = "window";
            string winName = head+to_string(i);
            cv::namedWindow(winName, WINDOW_NORMAL);
            int x = (channel[i] % showNum) * (fullScreenWidth/showNum);
            int y = (channel[i] / showNum) * ((fullScreenHeight-titleHeight)/showNum);
            cv::moveWindow(winName, x, y+titleHeight);
            cv::resizeWindow(winName, fullScreenWidth/showNum, (fullScreenHeight-titleHeight)/showNum);
            winNameList.push_back(winName);
        }
    }
    pthread_create(&exeTids, NULL, ModelExecute, (void*)&channelNum);
    pthread_create(&posTids, NULL, PostProcess, NULL);

    for(int i = 0; i < channelNum; i++) {
        pthread_detach(preTids[i]);
    }
    pthread_detach(exeTids);
    pthread_detach(posTids);
    
    while(!exitFlag) {
        sleep(10);
    }
    if (displayMode == "imshow") {
        cv::destroyAllWindows();
    }
    return 0;
}
